import numpy as np
import torch
import torch.nn.functional as F
from django.http import JsonResponse
from django.shortcuts import render
import base64
import cv2
import thePlugs.DBFace.common as common
from django_redis.serializers import json
from thePlugs.DBFace.model.DBFaceSmallH import DBFace
# Create your views here.

# 创建一个人脸检测的类
class DetectFace:
    def __init__(self):
        self.HAS_CUDA = torch.cuda.is_available()
        print(f"HAS_CUDA = {self.HAS_CUDA}")
        self.model = DBFace()
        self.model.eval()

    def nms(self,objs, iou=0.5):
        if objs is None or len(objs) <= 1:
            return objs

        objs = sorted(objs, key=lambda obj: obj.score, reverse=True)
        keep = []
        flags = [0] * len(objs)
        for index, obj in enumerate(objs):

            if flags[index] != 0:
                continue

            keep.append(obj)
            for j in range(index + 1, len(objs)):
                if flags[j] == 0 and obj.iou(objs[j]) > iou:
                    flags[j] = 1
        return keep

    # 检测人脸
    def detect(self,image, threshold=0.3, nms_iou=0.3):
        mean = [0.408, 0.447, 0.47]
        std = [0.289, 0.274, 0.278]

        image = common.pad(image)
        image = ((image / 255.0 - mean) / std).astype(np.float32)
        image = image.transpose(2, 0, 1)

        torch_image = torch.from_numpy(image)[None]
        if self.HAS_CUDA:
            torch_image = torch_image.cuda()

        hm, box, landmark = self.model(torch_image)
        hm_pool = F.max_pool2d(hm, 3, 1, 1)
        scores, indices = ((hm == hm_pool).float() * hm).view(1, -1).cpu().topk(1000)
        hm_height, hm_width = hm.shape[2:]

        scores = scores.squeeze()
        indices = indices.squeeze()
        ys = list((indices / hm_width).int().data.numpy())
        xs = list((indices % hm_width).int().data.numpy())
        scores = list(scores.data.numpy())
        box = box.cpu().squeeze().data.numpy()
        landmark = landmark.cpu().squeeze().data.numpy()

        stride = 4
        objs = []
        for cx, cy, score in zip(xs, ys, scores):
            if score < threshold:
                break

            x, y, r, b = box[:, cy, cx]
            xyrb = (np.array([cx, cy, cx, cy]) + [-x, -y, r, b]) * stride
            x5y5 = landmark[:, cy, cx]
            x5y5 = (common.exp(x5y5 * 4) + ([cx] * 5 + [cy] * 5)) * stride
            box_landmark = list(zip(x5y5[:5], x5y5[5:]))
            objs.append(common.BBox(0, xyrb=xyrb, score=score, landmark=box_landmark))
        return self.nms(objs, iou=nms_iou)

HAS_CUDA = torch.cuda.is_available()
print(f"HAS_CUDA = {HAS_CUDA}")

# 定义全局变量 摄像头
global cap,dbface,is_open_dbface
print("正在加载模型")
dbface = DetectFace()
cap = None
is_open_dbface = False


def img_to_base64(img):
    img_base64 = base64.b64encode(cv2.imencode('.jpg', img)[1]).decode()
    return str(img_base64)


def open_cap(request):
    msg=""
    global cap
    if cap==None:
        print("正在开始加载摄像头")
        cap = cv2.VideoCapture(0)
        cap.set(cv2.CAP_PROP_FRAME_WIDTH, 600)
        cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 400)
        msg="摄像头加载成功"
    else:
        msg = "摄像头已经加载"

    return JsonResponse({"status": 200, "msg": msg, "data":None})

def close_cap(request):
    msg=""
    global cap
    if cap==None:
        msg="摄像头已经关闭"
    else:
        cap.release()
        cv2.destroyAllWindows()
        cap=None
        msg = "摄像头关闭成功"

    return JsonResponse({"status": 200, "msg": msg, "data": None})

def send_picture(request):
    global cap,is_open_dbface
    if cap!=None:
        ok, frame = cap.read()
        if is_open_dbface:
            objs = dbface.detect(frame)
            for obj in objs:
                common.drawbbox(frame, obj)

        return JsonResponse({"status": 200, "msg": "OK", "image": img_to_base64(frame), "is_cap_closed": '0'})

    else:
        img = cv2.imread(r"E:\gitee\sp\myprject_202104\static\imgs\camera_is_closed.png")
        img_base64 = base64.b64encode(cv2.imencode('.png',img)[1]).decode()
        data = str(img_base64)
        return JsonResponse({"status": 200, "msg": "OK","image": data,"is_cap_closed":'1'})


def open_close_dbface(request):
    global is_open_dbface
    msg = ""
    if is_open_dbface:
        msg="人脸检测已经关闭"
    else:
        msg="人脸检测已开启"

    is_open_dbface = not is_open_dbface

    return JsonResponse({"status": 200, "msg": msg})



